import paramiko
import re
from time import sleep


# 定义一个类，表示一台远端算力主机
class Compute_Host(object):
    # 通过IP, 用户名，密码，超时时间初始化一个远程Linux主机
    def __init__(self, ip, username, password, timeout=30):
        self.ip = ip
        self.username = username
        self.password = password
        self.timeout = timeout
        # transport和chanel
        self.t = ''
        self.chan = ''
        # 链接失败的重试次数
        self.try_times = 3

    # 调用该方法连接远程主机
    def connect(self):
        while True:
            # 连接过程中可能会抛出异常，比如网络不通、链接超时
            try:
                self.t = paramiko.Transport(sock=(self.ip, 22))
                self.t.connect(username=self.username, password=self.password)
                self.chan = self.t.open_session()
                self.chan.settimeout(self.timeout)
                self.chan.get_pty()
                self.chan.invoke_shell()
                # 如果没有抛出异常说明连接成功，直接返回
                print(u'连接%s成功' % self.ip)
                # 接收到的网络数据解码为str
                print(self.chan.recv(65535).decode('utf-8'))
                return
            # 这里不对可能的异常如socket.error, socket.timeout细化，直接一网打尽
            except Exception:
                if self.try_times != 0:
                    print(u'连接%s失败，进行重试' % self.ip)
                    self.try_times -= 1
                else:
                    print(u'重试3次失败，结束程序')
                    exit(1)

    # 断开连接
    def close(self):
        self.chan.close()

        self.t.close()

    # 发送要执行的命令
    def send(self, cmd, pattern='~'):
        cmd += '\r'
        # 通过命令执行提示符来判断命令是否执行完成
        patt = pattern
        p = re.compile(patt)
        result = ''
        # 发送要执行的命令
        self.chan.send(cmd)
        # 回显很长的命令可能执行较久，通过循环分批次取回回显
        while True:
            sleep(0.5)
            ret = self.chan.recv(65535)
            ret = ret.decode('utf-8')
            result += ret
            if p.search(ret):
                print(result)
                return result

class Task:
    '''该类为整体任务的封装'''
    def __init__(self, job_id, send_command, inp_file):
        self.job_id = job_id
        self.send_command = send_command #这是发送命令的接口，由连接服务器的对象传入
        self.status = '队列中'
        self.computation_mode = '隐式' #隐式或者显式，默认隐式计算
        self.CPUs = 6 #默认使用6个CPU计算，可更改属性
        self.inp_file = inp_file #这个是输入文件地址
        self.GPUs = 0 #默认不使用GPU计算，可更改属性
        self.inp_name = '' #这个是输入文件名

    def run(self):
        if self.computation_mode == '显式':
            xianshi_command = f'abaqus job={self.inp_name} cpus={self.CPUs}'
            self.send_command(xianshi_command)
            sleep(1)
            self.status = '作业显式计算中'

        if self.computation_mode == '隐式':
            yinshi_command = f'abaqus job={self.inp_name} cpus={self.CPUs}'
            if self.GPUs != 0:
                yinshi_command += f'abaqus job={self.inp_name} cpus={self.CPUs} gpus={self.GPUs}'
            self.send_command(yinshi_command)
            sleep(1)
            self.status = '作业隐式计算中'

    def stop(self):
        self.send_command(f'abaqus terminate job={self.inp_name}')
        self.status = '作业终止'


class API_of_abaqus:
    def __init__(self):
        pass


class ResourceManager:
    '''该类为计算服务器的资源管理器'''
    def __init__(self):
        pass



class TaskScheduler:
    '''该类为任务调度器'''
    def __init__(self):
        pass


class Task_Database:
    '''该类为任务数据库'''
    def __init__(self):
        pass

if __name__ == '__main__':
    host = Compute_Host('192.168.0.183', 'jufeng', 'dagesofu')
    host.connect()
    # host.send('su root',r'.*(]#|密码).*')
    # host.send('123456')
    host.send('ls')
    host.close()

